import pygame
import random
import math
import sys
from pygame.locals import *

# 初始化pygame
pygame.init()

# 游戏常量
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
GRID_SIZE = 60
GRID_WIDTH = SCREEN_WIDTH // GRID_SIZE
GRID_HEIGHT = (SCREEN_HEIGHT - 100) // GRID_SIZE

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
YELLOW = (255, 255, 0)
BROWN = (139, 69, 19)
GRAY = (200, 200, 200)

# 创建屏幕
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("植物大战僵尸")

# 时钟
clock = pygame.time.Clock()
FPS = 30

# 字体
font = pygame.font.SysFont("SimHei", 24)
small_font = pygame.font.SysFont("SimHei", 16)

class GameObject:
    """游戏对象基类"""
    def __init__(self, x, y, width, height):
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.rect = pygame.Rect(x, y, width, height)
        
    def update(self):
        self.rect.x = self.x
        self.rect.y = self.y
        
    def draw(self, surface):
        pygame.draw.rect(surface, BLACK, self.rect)

class Plant(GameObject):
    """植物基类"""
    def __init__(self, x, y, cost, health, attack_power=0, attack_range=0):
        super().__init__(x, y, GRID_SIZE - 10, GRID_SIZE - 10)
        self.cost = cost
        self.health = health
        self.max_health = health
        self.attack_power = attack_power
        self.attack_range = attack_range
        self.attack_cooldown = 30  # 攻击冷却时间（帧）
        self.attack_timer = 0
        
    def update(self):
        super().update()
        if self.attack_timer > 0:
            self.attack_timer -= 1
            
    def take_damage(self, damage):
        self.health -= damage
        return self.health <= 0  # 返回是否死亡
        
    def draw_health_bar(self, surface):
        # 绘制生命值条
        bar_width = self.width
        bar_height = 5
        health_ratio = self.health / self.max_health
        
        # 背景条
        pygame.draw.rect(surface, RED, (self.x, self.y - 10, bar_width, bar_height))
        # 生命值条
        pygame.draw.rect(surface, GREEN, (self.x, self.y - 10, bar_width * health_ratio, bar_height))

class Sunflower(Plant):
    """向日葵：产生阳光"""
    def __init__(self, x, y):
        super().__init__(x, y, 50, 100)
        self.sun_timer = 120  # 每120帧产生一个阳光
        
    def update(self):
        super().update()
        self.sun_timer -= 1
        if self.sun_timer <= 0:
            self.sun_timer = 120
            return Sun(self.x + GRID_SIZE//2, self.y, 25)  # 产生阳光
        return None
        
    def draw(self, surface):
        pygame.draw.rect(surface, YELLOW, self.rect, border_radius=10)
        text = small_font.render("向", True, BLACK)
        surface.blit(text, (self.x + 10, self.y + 10))
        self.draw_health_bar(surface)

class Peashooter(Plant):
    """豌豆射手：发射豌豆攻击僵尸"""
    def __init__(self, x, y):
        super().__init__(x, y, 100, 100, 20, SCREEN_WIDTH)
        
    def attack(self, zombies):
        if self.attack_timer == 0:
            # 检查是否有僵尸在攻击范围内
            for zombie in zombies:
                if (zombie.y // GRID_SIZE == self.y // GRID_SIZE and 
                    zombie.x > self.x):
                    self.attack_timer = self.attack_cooldown
                    return Pea(self.x + self.width, self.y + self.height//2)
        return None
        
    def draw(self, surface):
        pygame.draw.rect(surface, GREEN, self.rect, border_radius=5)
        text = small_font.render("豌", True, BLACK)
        surface.blit(text, (self.x + 10, self.y + 10))
        self.draw_health_bar(surface)

class Pea(GameObject):
    """豌豆子弹"""
    def __init__(self, x, y):
        super().__init__(x, y, 10, 5)
        self.speed = 5
        self.damage = 20
        
    def update(self):
        super().update()
        self.x += self.speed
        # 超出屏幕则移除
        if self.x > SCREEN_WIDTH:
            return True
        return False
        
    def draw(self, surface):
        pygame.draw.rect(surface, GREEN, self.rect)

class Sun(GameObject):
    """阳光：收集后增加阳光值"""
    def __init__(self, x, y, value):
        super().__init__(x, y, 20, 20)
        self.value = value
        self.fall_speed = 1
        self.life_timer = 300  # 存在时间（帧）
        self.initial_y = y  # 记录初始y坐标
        # 下落1-2格后停止（每格高度=GRID_SIZE）
        self.max_fall_distance = GRID_SIZE * (1 + random.randint(0, 1))
        self.has_stopped = False  # 是否已停止下落
        
    def update(self):
        """更新位置和生命周期，返回是否需要移除"""
        super().update()
        
        # 未停止且未达到最大下落距离时继续下落
        if not self.has_stopped and (self.y - self.initial_y) < self.max_fall_distance:
            self.y += self.fall_speed
        else:
            self.has_stopped = True  # 标记为已停止
            
        self.life_timer -= 1
        return self.life_timer <= 0 or self.y > SCREEN_HEIGHT
        
    def draw(self, surface):
        pygame.draw.circle(surface, YELLOW, (self.x + 10, self.y + 10), 10)

class Zombie(GameObject):
    """僵尸基类"""
    def __init__(self, x, y, health=200, speed=1, damage=10):
        super().__init__(x, y, GRID_SIZE - 10, GRID_SIZE - 10)
        self.health = health
        self.max_health = health
        self.speed = speed
        self.damage = damage
        self.attack_cooldown = 30
        self.attack_timer = 0
        
    def update(self, plants):
        super().update()
        # 检查是否可以攻击植物
        attacking_plant = None
        for plant in plants:
            if (plant.y // GRID_SIZE == self.y // GRID_SIZE and 
                self.x < plant.x + plant.width and 
                self.x + self.width > plant.x):
                attacking_plant = plant
                break
                
        if attacking_plant:
            # 攻击植物
            if self.attack_timer == 0:
                attacking_plant.take_damage(self.damage)
                self.attack_timer = self.attack_cooldown
            else:
                self.attack_timer -= 1
        else:
            # 移动
            self.x -= self.speed
            
        # 检查是否到达终点
        return self.x < -self.width
        
    def take_damage(self, damage):
        self.health -= damage
        return self.health <= 0
        
    def draw_health_bar(self, surface):
        bar_width = self.width
        bar_height = 5
        health_ratio = self.health / self.max_health
        
        pygame.draw.rect(surface, RED, (self.x, self.y - 10, bar_width, bar_height))
        pygame.draw.rect(surface, GREEN, (self.x, self.y - 10, bar_width * health_ratio, bar_height))
        
    def draw(self, surface):
        pygame.draw.rect(surface, BROWN, self.rect)
        text = small_font.render("僵", True, BLACK)
        surface.blit(text, (self.x + 10, self.y + 10))
        self.draw_health_bar(surface)

class Game:
    def __init__(self):
        self.plants = []
        self.zombies = []
        self.peas = []
        self.suns = []
        self.sun_points = 100
        self.selected_plant = None
        self.game_over = False
        self.zombie_spawn_timer = 180  # 僵尸生成计时器
        self.wave = 1
        
    def handle_events(self):
        for event in pygame.event.get():
            if event.type == QUIT:
                pygame.quit()
                sys.exit()
                
            if event.type == MOUSEBUTTONDOWN and not self.game_over:
                x, y = pygame.mouse.get_pos()
                
                # 收集阳光
                for sun in self.suns[:]:
                    if sun.rect.collidepoint(x, y):
                        self.sun_points += sun.value
                        self.suns.remove(sun)
                        
                # 选择植物
                if y > SCREEN_HEIGHT - 100:
                    # 向日葵
                    if 50 <= x <= 50 + GRID_SIZE and self.sun_points >= 50:
                        self.selected_plant = "sunflower"
                    # 豌豆射手
                    elif 120 <= x <= 120 + GRID_SIZE and self.sun_points >= 100:
                        self.selected_plant = "peashooter"
                # 种植植物
                else:
                    if self.selected_plant and self.can_plant(x, y):
                        grid_x = (x // GRID_SIZE) * GRID_SIZE
                        grid_y = (y // GRID_SIZE) * GRID_SIZE
                        
                        if self.selected_plant == "sunflower" and self.sun_points >= 50:
                            self.plants.append(Sunflower(grid_x, grid_y))
                            self.sun_points -= 50
                        elif self.selected_plant == "peashooter" and self.sun_points >= 100:
                            self.plants.append(Peashooter(grid_x, grid_y))
                            self.sun_points -= 100
                            
                        self.selected_plant = None
    
    def can_plant(self, x, y):
        grid_x = x // GRID_SIZE
        grid_y = y // GRID_SIZE
        
        # 检查是否在可种植区域
        if grid_y >= GRID_HEIGHT:
            return False
            
        # 检查是否已有植物
        for plant in self.plants:
            if (plant.x // GRID_SIZE == grid_x and 
                plant.y // GRID_SIZE == grid_y):
                return False
                
        return True
    
    def update(self):
        if self.game_over:
            return
            
        # 更新植物
        for plant in self.plants[:]:
            plant.update()
            # 向日葵产生阳光
            if isinstance(plant, Sunflower):
                sun = plant.update()
                if sun:
                    self.suns.append(sun)
            # 豌豆射手攻击
            elif isinstance(plant, Peashooter):
                pea = plant.attack(self.zombies)
                if pea:
                    self.peas.append(pea)
        
        # 更新豌豆
        for pea in self.peas[:]:
            if pea.update():
                self.peas.remove(pea)
            else:
                # 检测是否击中僵尸
                for zombie in self.zombies[:]:
                    if pea.rect.colliderect(zombie.rect):
                        if zombie.take_damage(pea.damage):
                            self.zombies.remove(zombie)
                        self.peas.remove(pea)
                        break
        
        # 更新阳光
        for sun in self.suns[:]:
            if sun.update():
                self.suns.remove(sun)
        
        # 生成僵尸
        self.zombie_spawn_timer -= 1
        if self.zombie_spawn_timer <= 0:
            row = random.randint(0, GRID_HEIGHT - 1)
            self.zombies.append(Zombie(SCREEN_WIDTH, row * GRID_SIZE))
            # 随着波数增加，僵尸生成速度加快
            self.zombie_spawn_timer = max(60, 180 - (self.wave - 1) * 10)
            # 每生成10个僵尸，波数增加
            if len(self.zombies) % 10 == 0:
                self.wave += 1
        
        # 更新僵尸
        for zombie in self.zombies[:]:
            if zombie.update(self.plants):
                # 僵尸到达终点，游戏结束
                self.game_over = True
            # 检查僵尸是否死亡
            if zombie.health <= 0:
                self.zombies.remove(zombie)
        
        # 移除死亡的植物
        for plant in self.plants[:]:
            if plant.health <= 0:
                self.plants.remove(plant)
    
    def draw(self):
        screen.fill(GRAY)
        
        # 绘制网格
        for x in range(0, SCREEN_WIDTH, GRID_SIZE):
            pygame.draw.line(screen, WHITE, (x, 0), (x, SCREEN_HEIGHT - 100), 1)
        for y in range(0, SCREEN_HEIGHT - 100, GRID_SIZE):
            pygame.draw.line(screen, WHITE, (0, y), (SCREEN_WIDTH, y), 1)
        
        # 绘制植物
        for plant in self.plants:
            plant.draw(screen)
        
        # 绘制豌豆
        for pea in self.peas:
            pea.draw(screen)
        
        # 绘制阳光
        for sun in self.suns:
            sun.draw(screen)
        
        # 绘制僵尸
        for zombie in self.zombies:
            zombie.draw(screen)
        
        # 绘制工具栏
        pygame.draw.rect(screen, BROWN, (0, SCREEN_HEIGHT - 100, SCREEN_WIDTH, 100))
        
        # 绘制植物卡片
        # 向日葵
        pygame.draw.rect(screen, YELLOW, (50, SCREEN_HEIGHT - 90, GRID_SIZE, GRID_SIZE), border_radius=10)
        text = small_font.render("向日葵 (50)", True, BLACK)
        screen.blit(text, (50, SCREEN_HEIGHT - 30))
        
        # 豌豆射手
        pygame.draw.rect(screen, GREEN, (120, SCREEN_HEIGHT - 90, GRID_SIZE, GRID_SIZE), border_radius=5)
        text = small_font.render("豌豆射手 (100)", True, BLACK)
        screen.blit(text, (120, SCREEN_HEIGHT - 30))
        
        # 绘制阳光值
        sun_text = font.render(f"阳光: {self.sun_points}", True, YELLOW)
        screen.blit(sun_text, (SCREEN_WIDTH - 150, SCREEN_HEIGHT - 50))
        
        # 绘制当前波数
        wave_text = font.render(f"波数: {self.wave}", True, WHITE)
        screen.blit(wave_text, (SCREEN_WIDTH - 300, SCREEN_HEIGHT - 50))
        
        # 绘制选中的植物提示
        if self.selected_plant:
            提示 = "向日葵" if self.selected_plant == "sunflower" else "豌豆射手"
            hint_text = font.render(f"已选择: {提示}", True, WHITE)
            screen.blit(hint_text, (300, SCREEN_HEIGHT - 50))
        
        # 游戏结束
        if self.game_over:
            overlay = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT), pygame.SRCALPHA)
            overlay.fill((0, 0, 0, 180))
            screen.blit(overlay, (0, 0))
            
            game_over_text = font.render("游戏结束!", True, RED)
            restart_text = font.render("按任意键重新开始", True, WHITE)
            
            screen.blit(game_over_text, (SCREEN_WIDTH//2 - 80, SCREEN_HEIGHT//2 - 30))
            screen.blit(restart_text, (SCREEN_WIDTH//2 - 120, SCREEN_HEIGHT//2 + 30))
    
    def reset(self):
        """重置游戏"""
        self.__init__()

def main():
    game = Game()
    
    while True:
        game.handle_events()
        
        # 检查游戏是否结束并重新开始
        if game.game_over:
            keys = pygame.key.get_pressed()
            if any(keys):
                game.reset()
        
        game.update()
        game.draw()
        
        pygame.display.flip()
        clock.tick(FPS)

if __name__ == "__main__":
    main()